from lida import Manager,TextGenerationConfig,llm
from lida.datamodel import Goal 
from lida.utils import plot_raster 

#from llmx import llm
model_name = "moonshot-v1-8k"
library = "seaborn"

text_gen = llm(provider="openai", api_base="https://api.moonshot.cn/v1", api_key=API_KEY, model=model_name)

textgen_config = TextGenerationConfig(n=1, temperature=0.5, model=model_name, use_cache=True)


lida = Manager(text_gen)

summary = lida.summarize("../data/cars.csv", summary_method="default", textgen_config=textgen_config)  
print(summary)
# datamodel 总共有4 个object，分别是index, question, visualization and rationale
custom_goal = Goal( 
    index= 0 , 
    question= "What is the distribution of the Type?" , 
    visualization= "Bar Chart" , 
    rationale= "The type of the car is an important feature of the dataset."
)
# 生成图表
custom_chart = lida.visualize(summary=summary, goal=custom_goal, textgen_config=textgen_config , library=library) 
custom_chart[ 0 ].savefig('example_18_custom.png') # 保存图表
# 编辑客制化生成图表
custom_instructions = [ "change the color to red tone on tone color" ] # 改变Bar Chart 的颜色
edited_custom_charts = lida.edit(code= custom_chart[ 0 ].code, summary=summary, instructions=custom_instructions) 
custom_chart[ 0 ].savefig('example_18_custom_edited.png') # 保存编辑后的图表